I think it would make more sense to return the logprob as part of the step method?

Then we could create a custom object that contains both the text output and the logprob. This might do for now, give me some time to think about it.

What are you using this for?

I think it would make more sense to return the logprob as part of the step method?

Then we could create a custom object that contains both the text output and the logprob. This might do for now, give me some time to think about it.

If you include it in the step are you thinking to return the logit only for the sampled token or all of the logits. I think it would be useful for the user (at least me 😄 ) to have access to all of the logits if needed.

What are you using this for?

It's just very useful to be able to analyse the next token logits to e.g. assess the confidence of the model on a True/False question for various questions (and this pairs very well with regex).

Ah I see what you mean, step already outputs the probs as an additional output.

It would be useful to have an additional method with a simple API like
generate.continuation(model).next_token_logits(prompt)
or
generate.continuation(model).next_token_probs(prompt)

I.e. returns the next token logits or probs without actually doing any sampling.

Ok I've thought about this a bit more.

IMO the best API would be to have a .next_token_logits(prompt) or .next_token_probs(prompt) function that allows the user to examine the distribution without sampling. I think it's good to have this separate from the step function which samples and appends the new token, which is a different use case.

With this new function it may also make sense to have step only return the token_ids and not probs. This could also avoid some confusion as the token_ids relate to the whole sequence and probs only the appended token.

Was there a specific reason to have step return probs as well?

WDYT? @rlouf @brandonwillard

Was there a specific reason to have step return probs as well?

Unnecessary anticipation of SMC sampling.

WDYT? @rlouf @brandonwillard

Sounds good to me.

Do you prefer .logits or .probs or .next_token_logits or .next_token_probs?

I can imagine use cases where the user wants both logits or probs so IMO best to go with .logits or next_token_logits as this has (slightly) less computation and the user can easily call softmax if they need, this is also more inline with the output from self.model.

next_token_logits sounds good

I'm now thinking that indeed the probs might be useful for SMC 😃

Let's leave step returning probs for now

I've had a look at this and I'm not sure it's needed.

The pros:

• It gives finer scale control for people inheriting sequence
• It allows people to avoid superfluous sampling if they're just interested in the logits (eg speculative decoding)
• It avoids needing to torch.log the probabilities, which might be useful for numerical stability.

The cons:

• There's a real risk that people will use this and step, which would lead to an extra model call.
• The information that is needed is already present in the output of step modulo a logarithm.

@rlouf - Is it possible to consistently work with logits rather than probs, just for numerical stability. I've had a look at the code and it's not clear to me why you're not using torch.multinomial

(Also sorry @rlouf I must've mis-clicked. I don't think I meant to remove you as a reviewer)

Solved in #366